{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Backprop demo on a single unit linear MLP\n",
    "This is a demonstration of how backpropagation works on a single unit MLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A single neuron/unit MLP with linear activation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(1, 1)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build a network. Model a function: y = 2*x + 1\n",
    "\n",
    "Perform supervised learning using the following dataset:\n",
    "\n",
    "\n",
    "| Step |  x   |  y  |\n",
    "| :-- | ---: | ---:|\n",
    "|  0   | 0.   |  1. |\n",
    "|  1   | 1.   |  3. |\n",
    "\n",
    "At each step, perform backprop and print the gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------------0-----------------------\n",
      "0. Bias grad:  None\n",
      "0. Weights grad:  None\n",
      "0. Bias:  tensor(0., grad_fn=<SelectBackward>)\n",
      "0. Weights:  tensor([1.], grad_fn=<SelectBackward>)\n",
      "0. Input:  tensor([0.])\n",
      "0. Predicted Output:  tensor([0.], grad_fn=<AddBackward0>)\n",
      "0. Loss:  tensor(1., grad_fn=<MseLossBackward>)\n",
      "\n",
      "---------------------1-----------------------\n",
      "1. Bias grad:  tensor([-2.])\n",
      "1. Weights grad:  tensor([[0.]])\n",
      "1. Bias:  tensor(0.2000, grad_fn=<SelectBackward>)\n",
      "1. Weights:  tensor([1.], grad_fn=<SelectBackward>)\n",
      "1. Input:  tensor([1.])\n",
      "1. Predicted Output:  tensor([1.2000], grad_fn=<AddBackward0>)\n",
      "1. Loss:  tensor(3.2400, grad_fn=<MseLossBackward>)\n",
      "\n",
      "---------------------2-----------------------\n",
      "2. Bias grad:  tensor([-3.6000])\n",
      "2. Weights grad:  tensor([[-3.6000]])\n",
      "2. Bias:  tensor(0.5600, grad_fn=<SelectBackward>)\n",
      "2. Weights:  tensor([1.3600], grad_fn=<SelectBackward>)\n"
     ]
    }
   ],
   "source": [
    "net = Net()\n",
    "\n",
    "# for ease of backprop demo, we set weight=1.0 and bias=0.0\n",
    "net.fc1.bias = torch.nn.Parameter(torch.tensor([0.]))\n",
    "net.fc1.weight = torch.nn.Parameter(torch.tensor([[1.]]))\n",
    "print(\"---------------------0-----------------------\")\n",
    "print(\"0. Bias grad: \", net.fc1.bias.grad)\n",
    "print(\"0. Weights grad: \", net.fc1.weight.grad)\n",
    "print(\"0. Bias: \", net.fc1.bias[0])\n",
    "print(\"0. Weights: \", net.fc1.weight[0])\n",
    "\n",
    "# x=0, y=1.\n",
    "input = torch.tensor([0.])\n",
    "print(\"0. Input: \", input)\n",
    "input.unsqueeze(0)\n",
    "output = net(input)\n",
    "print(\"0. Predicted Output: \", output)\n",
    "target = torch.tensor([1.])\n",
    "target = target.view(1, -1)\n",
    "\n",
    "# Use MSE Loss\n",
    "criterion = nn.MSELoss()\n",
    "loss = criterion(output, target)\n",
    "print(\"0. Loss: \", loss)\n",
    "\n",
    "# Use SGD optimizer with learning rate of 0.1\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.1)\n",
    "# Clear optimizer gradient\n",
    "optimizer.zero_grad()\n",
    "# Perform backprop\n",
    "loss.backward(retain_graph=True)\n",
    "# Update weight and bias\n",
    "optimizer.step()\n",
    "print(\"\\n---------------------1-----------------------\")\n",
    "print(\"1. Bias grad: \", net.fc1.bias.grad)\n",
    "print(\"1. Weights grad: \",net.fc1.weight.grad)\n",
    "print(\"1. Bias: \", net.fc1.bias[0])\n",
    "print(\"1. Weights: \",net.fc1.weight[0])\n",
    "\n",
    "# x=1.0, y=3.0\n",
    "input = torch.tensor([1.])\n",
    "print(\"1. Input: \", input)\n",
    "input.unsqueeze(0)\n",
    "output = net(input)\n",
    "print(\"1. Predicted Output: \", output)\n",
    "target = torch.tensor([3.])\n",
    "target = target.view(1, -1)\n",
    "\n",
    "optimizer.zero_grad()\n",
    "loss = criterion(output, target)\n",
    "print(\"1. Loss: \", loss)\n",
    "loss.backward()\n",
    "optimizer.step()\n",
    "print(\"\\n---------------------2-----------------------\")\n",
    "print(\"2. Bias grad: \", net.fc1.bias.grad)\n",
    "print(\"2. Weights grad: \",net.fc1.weight.grad)\n",
    "print(\"2. Bias: \", net.fc1.bias[0])\n",
    "print(\"2. Weights: \",net.fc1.weight[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
